#
# Copyright (C) [2024] Xingyun Integrated Circuit, Inc.
#
#   GreenCode was a private technology asset of Xingyun Integrated Circuit， Inc （Confidential）
#   Author :  Shawn.Tan
#   Date : 2025.10.28
#
#   History : Initial Version 2025.10.28
#
#
import UtilityFunctions
from Enums import EStateElementType, EStateTransitionType

from base.StateTransitionHandler import StateTransitionHandler
from gpgpu.Utils import LoadGPR64


#  This class generates instructions to update the system State according to
#  the MemoryStateElements provided to the processStateElement() and
# processStateElements() methods.
class MemoryStateTransitionHandlerGPGPU(StateTransitionHandler):
    def __init__(self, aGenThread):
        super().__init__(aGenThread)

        self._mHelperGprSet = StateTransitionHelperGprSet(self)

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.Memory:
            return False

        (
            base_reg_index,
            mem_val_reg_index,
        ) = self._mHelperGprSet.acquireHelperGprs(2)

        load_gpr64_seq = LoadGPR64(self.genThread)
        load_gpr64_seq.load(base_reg_index, aStateElem.getStartAddress())
        load_gpr64_seq.load(mem_val_reg_index, aStateElem.getValues()[0])

        self._genStoreInstruction(base_reg_index, mem_val_reg_index, 0)

        self._mHelperGprSet.releaseHelperGprs()

        return True

    # Execute the State changes represented by the StateElements. Only
    # instances of the StateElement types for which the StateTransitionHandler
    # has been registered will be passed to this method. Other
    # StateTransitionHandlers will process the other StateElement types. It is
    # important to avoid making changes to entities represented by
    # StateElements that have already been processed. Changes to entities
    # represented by StateElements that will be processed later are permitted.
    #
    #  @param aStateElems A list of all StateElement objects of a particular
    #       type.
    def processStateElements(self, aStateElems):
        if aStateElems[0].getStateElementType() != EStateElementType.Memory:
            self.error(
                "This StateTransitionHandler can only process "
                "StateElements of type %s" % EStateElementType.Memory
            )

        (
            base_reg_index,
            mem_val_reg_index,
        ) = self._mHelperGprSet.acquireHelperGprs(2)

        load_gpr64_seq = LoadGPR64(self.genThread)
        offset_limit = 2 ** 11
        base_addr = -offset_limit
        for state_elem in aStateElems:
            load_gpr64_seq.load(mem_val_reg_index, state_elem.getValues()[0])

            # This is a minor optimization for the likely case in which the
            # starting addresses of consecutive State Elements are close
            # together. Instead of loading the base register from scratch, we
            # can offset from its current value.
            offset = state_elem.getStartAddress() - base_addr
            if (offset < -offset_limit) or (offset >= offset_limit):
                base_addr = state_elem.getStartAddress()
                load_gpr64_seq.load(base_reg_index, base_addr)
                offset = 0

            self._genStoreInstruction(base_reg_index, mem_val_reg_index, offset)

        self._mHelperGprSet.releaseHelperGprs()

    # Generate a store instruction.
    #
    #  @param aBaseRegIndex The base address register.
    #  @param aMemValRegIndex The memory value source register.
    #  @param aOffset The address offset value.
    def _genStoreInstruction(self, aBaseRegIndex, aSourceRegIndex, aOffset):
        if self.getGlobalState("AppRegisterWidth") == 32:
            self.genInstruction(
                "SW##CALM",
                {
                    "rs1": aBaseRegIndex,
                    "rs2": aSourceRegIndex,
                    "simm12": aOffset,
                    "NoRestriction": 1,
                },
            )
        else:
            self.genInstruction(
                "SD##CALM",
                {
                    "rs1": aBaseRegIndex,
                    "rs2": aSourceRegIndex,
                    "simm12": aOffset,
                    "NoRestriction": 1,
                },
            )


#  This class generates instructions to update the system State according to
#  the vector register StateElements provided to the processStateElement() and
# processStateElements() methods.
class VectorRegisterStateTransitionHandlerGPGPU(StateTransitionHandler):
    def __init__(self, aGenThread):
        super().__init__(aGenThread)

        self._mHelperGprSet = StateTransitionHelperGprSet(self)

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.VectorRegister:
            return False

        (mem_block_ptr_index,) = self._mHelperGprSet.acquireHelperGprs(1)

        self.initializeMemoryBlock(mem_block_ptr_index, (aStateElem,))
        self.genInstruction(
            "VL1RE8.V##CALM",
            {
                "vd": aStateElem.getRegisterIndex(),
                "rs1": mem_block_ptr_index,
                "NoRestriction": 1,
            },
        )

        self._mHelperGprSet.releaseHelperGprs()

        return True

    # Execute the State changes represented by the StateElements. Only
    # instances of the StateElement types for which the StateTransitionHandler
    # has been registered will be passed to this method. Other
    # StateTransitionHandlers will process the other StateElement types. It is
    # important to avoid making changes to entities represented by
    # StateElements that have already been processed. Changes to entities
    # represented by StateElements that will be processed later are permitted.
    #
    #  @param aStateElems A list of all StateElement objects of a particular
    #       type.
    def processStateElements(self, aStateElems):
        if aStateElems[0].getStateElementType() != EStateElementType.VectorRegister:
            self.error(
                "This StateTransitionHandler can only process "
                "StateElements of type %s" % EStateElementType.VectorRegister
            )

        (mem_block_ptr_index,) = self._mHelperGprSet.acquireHelperGprs(1)
        self.initializeMemoryBlock(mem_block_ptr_index, aStateElems)

        for state_elem in aStateElems:
            self.genInstruction(
                "VL1RE8.V##CALM",
                {
                    "vd": state_elem.getRegisterIndex(),
                    "rs1": mem_block_ptr_index,
                    "NoRestriction": 1,
                },
            )
            offset = len(state_elem.getValues()) * 8
            self.genInstruction(
                "ADDI##CALM",
                {
                    "rd": mem_block_ptr_index,
                    "rs1": mem_block_ptr_index,
                    "simm12": offset,
                },
            )

        self._mHelperGprSet.releaseHelperGprs()


# This class generates instructions to update the system State according to
# the system register StateElements provided to the processStateElement() and
# processStateElements() methods.
class SystemRegisterStateTransitionHandlerGPGPU(StateTransitionHandler):
    def __init__(self, aGenThread):
        super().__init__(aGenThread)

        self._mHelperGprSet = StateTransitionHelperGprSet(self)

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.SystemRegister:
            return False

        (
            reg_val_gpr_index,
            scratch_gpr_index,
        ) = self._mHelperGprSet.acquireHelperGprs(2, aValidate=False)

        load_gpr64_seq = LoadGPR64(self.genThread)
        load_gpr64_seq.load(reg_val_gpr_index, aStateElem.getValues()[0])
        self._processSystemRegisterStateElement(aStateElem, reg_val_gpr_index, scratch_gpr_index)

        self._mHelperGprSet.releaseHelperGprs()

        return True

    # Execute the State changes represented by the StateElements. Only
    # instances of the StateElement types for which the StateTransitionHandler
    # has been registered will be passed to this method. Other
    # StateTransitionHandlers will process the other StateElement types. It is
    # important to avoid making changes to entities represented by
    # StateElements that have already been processed. Changes to entities
    # represented by StateElements that will be processed later are permitted.
    #
    #  @param aStateElems A list of all StateElement objects of a particular
    #       type.
    def processStateElements(self, aStateElems):
        if aStateElems[0].getStateElementType() != EStateElementType.SystemRegister:
            self.error(
                "This StateTransitionHandler can only process "
                "StateElements of type %s" % EStateElementType.SystemRegister
            )

        (
            mem_block_ptr_index,
            reg_val_gpr_index,
            scratch_gpr_index,
        ) = self._mHelperGprSet.acquireHelperGprs(3)
        self.initializeMemoryBlock(mem_block_ptr_index, aStateElems)
        (mem_block_ptr_val, _) = self.readRegister("x%d" % mem_block_ptr_index)

        load_gpr64_seq = LoadGPR64(self.genThread)
        offset_limit = 2 ** 11
        offset = 0
        for state_elem in aStateElems:
            # This is a minor optimization that offsets from the memory block
            # pointer when possible rather than adjusting its value
            if offset >= offset_limit:
                mem_block_ptr_val += offset
                load_gpr64_seq.load(mem_block_ptr_index, mem_block_ptr_val)
                offset = 0

            self._genLoadInstruction(reg_val_gpr_index, mem_block_ptr_index, offset)
            self._processSystemRegisterStateElement(
                state_elem, reg_val_gpr_index, scratch_gpr_index
            )

            offset += 8

        self._mHelperGprSet.releaseHelperGprs()

    # Execute the State change represented by the system register StateElement.
    #
    #  @param aStateElem A system register StateElement object.
    #  @param aRegValGprIndex A GPR containing the system register value.
    #  @param aScratchGprIndex A GPR to be used as a scratch register.
    def _processSystemRegisterStateElement(self, aStateElem, aRegValGprIndex, aScratchGprIndex):
        if aStateElem.getName() == "vtype":
            self._processVtypeStateElement(aRegValGprIndex)
        elif aStateElem.getName() == "vl":
            self._processVlStateElement(aRegValGprIndex, aScratchGprIndex)
        else:
            self.genInstruction(
                "CSRRW#register#CALM",
                {
                    "rd": 0,
                    "rs1": aRegValGprIndex,
                    "csr": aStateElem.getRegisterIndex(),
                },
            )

    # Execute the State change represented by the vtype StateElement.
    #
    #  @param aRegValGprIndex A GPR containing the vtype register value.
    def _processVtypeStateElement(self, aRegValGprIndex):
        # Set rd and rs1 to x0 to preserve the value of vl
        self.genInstruction("VSETVL##CALM", {"rd": 0, "rs1": 0, "rs2": aRegValGprIndex})

    # Execute the State change represented by the vl StateElement.
    #
    #  @param aRegValGprIndex A GPR containing the vl register value.
    #  @param aScratchGprIndex A GPR to be used as a scratch register.
    def _processVlStateElement(self, aRegValGprIndex, aScratchGprIndex):
        # Read the current value of vtype and pass it to VSETVL in order to
        # preserve the value of vtype
        self.genInstruction(
            "CSRRS#register#CALM",
            {
                "rd": aScratchGprIndex,
                "rs1": 0,
                "csr": self.getRegisterIndex("vtype"),
            },
        )
        self.genInstruction(
            "VSETVL##CALM",
            {"rd": 0, "rs1": aRegValGprIndex, "rs2": aScratchGprIndex},
        )

    # Generate a load instruction.
    #
    #  @param aDestRegIndex The destination register.
    #  @param aBaseRegIndex The base address register.
    #  @param aOffset The address offset value.
    def _genLoadInstruction(self, aDestRegIndex, aBaseRegIndex, aOffset):
        if self.getGlobalState("AppRegisterWidth") == 32:
            self.genInstruction(
                "LW##CALM",
                {
                    "rd": aDestRegIndex,
                    "rs1": aBaseRegIndex,
                    "simm12": aOffset,
                    "NoRestriction": 1,
                },
            )
        else:
            self.genInstruction(
                "LD##CALM",
                {
                    "rd": aDestRegIndex,
                    "rs1": aBaseRegIndex,
                    "simm12": aOffset,
                    "NoRestriction": 1,
                },
            )


# This class generates instructions to update the system State according to
# the GPR StateElements provided to the processStateElement() and
# processStateElements() methods.
class GprStateTransitionHandlerGPGPU(StateTransitionHandler):

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.GPR:
            return False

        load_gpr64_seq = LoadGPR64(self.genThread)
        load_gpr64_seq.load(aStateElem.getRegisterIndex(), aStateElem.getValues()[0])
        return True

    # Execute the State changes represented by the StateElements. Only
    # instances of the StateElement types for which the StateTransitionHandler
    # has been registered will be passed to this method. Other
    # StateTransitionHandlers will process the other StateElement types. It is
    # important to avoid making changes to entities represented by
    # StateElements that have already been processed. Changes to entities
    # represented by StateElements that will be processed later are permitted.
    #
    #  @param aStateElems A list of all StateElement objects of a particular
    #       type.
    def processStateElements(self, aStateElems):
        if aStateElems[0].getStateElementType() != EStateElementType.GPR:
            self.error(
                "This StateTransitionHandler can only process "
                "StateElements of type %s" % EStateElementType.GPR
            )

        # The logic below assumes there will never be more than 2,048 GPR
        # StateElements
        if len(aStateElems) > 2048:
            self.error("Unexpected number of GPR StateElements: %d" % len(aStateElems))

        # Use the last GPR as the memory block pointer
        mem_block_ptr_index = aStateElems[-1].getRegisterIndex()
        self.initializeMemoryBlock(mem_block_ptr_index, aStateElems)

        ld_instr = "LW##CALM" if self.getGlobalState("AppRegisterWidth") == 32 else "LD##CALM"

        offset = 0
        for state_elem in aStateElems[:-1]:
            self.genInstruction(
                ld_instr,
                {
                    "rd": state_elem.getRegisterIndex(),
                    "rs1": mem_block_ptr_index,
                    "simm12": offset,
                    "NoRestriction": 1,
                },
            )
            offset += 8

        # Load the last StateElement value into the memory block pointer
        # register as the last step
        self.genInstruction(
            ld_instr,
            {
                "rd": mem_block_ptr_index,
                "rs1": mem_block_ptr_index,
                "simm12": offset,
                "NoRestriction": 1,
            },
        )


#  This class generates instructions to update the system State according to
#  the VM context StateElements provided to the processStateElement() and
# processStateElements() methods.
class VmContextStateTransitionHandlerGPGPU(StateTransitionHandler):
    def __init__(self, aGenThread):
        super().__init__(aGenThread)

        self._mHelperGprSet = StateTransitionHelperGprSet(self)

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.VmContext:
            return False

        (reg_val_gpr_index,) = self._mHelperGprSet.acquireHelperGprs(1)

        # Ensure the register is initialized before attempting to read it
        self.randomInitializeRegister(aStateElem.getRegisterName())
        (reg_val, _) = self.readRegister(aStateElem.getRegisterName())

        reg_val = combine_register_value_with_field_value(
            self,
            aStateElem.getRegisterName(),
            reg_val,
            aStateElem.getRegisterFieldName(),
            aStateElem.getValues()[0],
        )

        load_gpr64_seq = LoadGPR64(self.genThread)
        load_gpr64_seq.load(reg_val_gpr_index, reg_val)
        self.genInstruction(
            "CSRRW#register#CALM",
            {
                "rd": 0,
                "rs1": reg_val_gpr_index,
                "csr": self.getRegisterIndex(aStateElem.getRegisterName()),
            },
        )
        self.updateVm()

        self._mHelperGprSet.releaseHelperGprs()

        return True

    # Execute the State changes represented by the StateElements. Only
    # instances of the StateElement types for which the StateTransitionHandler
    # has been registered will be passed to this method. Other
    # StateTransitionHandlers will process the other StateElement types. It is
    # important to avoid making changes to entities represented by
    # StateElements that have already been processed. Changes to entities
    # represented by StateElements that will be processed later are permitted.
    #
    #  @param aStateElems A list of all StateElement objects of a particular
    #       type.
    def processStateElements(self, aStateElems):
        if aStateElems[0].getStateElementType() != EStateElementType.VmContext:
            self.error(
                "This StateTransitionHandler can only process "
                "StateElements of type %s" % EStateElementType.VmContext
            )

        (reg_val_gpr_index,) = self._mHelperGprSet.acquireHelperGprs(1)

        # Consolidate StateElements by register, so we only need to set each
        # register once
        reg_fields = self._groupStateElementFieldsByRegister(aStateElems)

        load_gpr64_seq = LoadGPR64(self.genThread)
        for (
            reg_name,
            (reg_field_names, reg_field_values),
        ) in reg_fields.items():
            # Ensure the register is initialized before attempting to read it
            self.randomInitializeRegister(reg_name)
            (reg_val, _) = self.readRegister(reg_name)

            for (i, reg_field_name) in enumerate(reg_field_names):
                reg_val = combine_register_value_with_field_value(
                    self,
                    reg_name,
                    reg_val,
                    reg_field_name,
                    reg_field_values[i],
                )

            load_gpr64_seq.load(reg_val_gpr_index, reg_val)
            self.genInstruction(
                "CSRRW#register#CALM",
                {
                    "rd": 0,
                    "rs1": reg_val_gpr_index,
                    "csr": self.getRegisterIndex(reg_name),
                },
            )
            self.updateVm()

        self._mHelperGprSet.releaseHelperGprs()

    # Create a dictionary of register names mapped to a tuple of a list of
    # register field names and a list of register field values.
    #
    #  @param aStateElems A list of all VmContextStateElement objects.
    def _groupStateElementFieldsByRegister(self, aStateElems):
        reg_fields = {}
        for state_elem in aStateElems:
            if state_elem.getRegisterName() in reg_fields:
                reg_fields_entry = reg_fields[state_elem.getRegisterName()]
                reg_fields_entry[0].append(state_elem.getRegisterFieldName())
                reg_fields_entry[1].append(state_elem.getValues()[0])
            else:
                reg_fields[state_elem.getRegisterName()] = (
                    [state_elem.getRegisterFieldName()],
                    [state_elem.getValues()[0]],
                )

        return reg_fields


#  This class generates instructions to update the system State according to
#  the privlege level StateElements provided to the processStateElement()
# method. processStateElements() is not overriden because only one privilege
# level StateElement can be specified per State.
class PrivilegeLevelStateTransitionHandlerGPGPU(StateTransitionHandler):
    def __init__(self, aGenThread):
        super().__init__(aGenThread)

        self._mHelperGprSet = StateTransitionHelperGprSet(self)

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.PrivilegeLevel:
            return False

        target_priv_level = aStateElem.getValues()[0]
        if target_priv_level != self.getPEstate("PrivilegeLevel"):
            self.systemCall({"PrivilegeLevel": target_priv_level})
        # Else we're already at the target privilege level, so there is nothing
        # to do

        return True


# This class generates instructions to update the system State according to
# the PC StateElements provided to the processStateElement() method.
# processStateElements() is not overriden because only one PC StateElement can
# be specified per State.
class PcStateTransitionHandlerGPGPU(StateTransitionHandler):
    def __init__(self, aGenThread):
        super().__init__(aGenThread)

        self._mHelperGprSet = StateTransitionHelperGprSet(self)

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.PC:
            return False

        target_addr = aStateElem.getValues()[0]

        # Try using a PC-relative branch first to limit the number of generated
        # instructions; fall back to a register branch if that fails.
        (branch_offset, is_valid, num_hw) = self.getBranchOffset(
            self.getPEstate("PC"), target_addr, 20, 1
        )
        if is_valid:
            self.genInstruction(
                "JAL##CALM",
                {"rd": 0, "simm20": branch_offset, "NoRestriction": 1},
            )
        else:
            (branch_gpr_index,) = self._mHelperGprSet.acquireHelperGprs(1)

            load_gpr = LoadGPR64(self.genThread)
            load_gpr.load(branch_gpr_index, target_addr)
            self.genInstruction(
                "JALR##CALM",
                {
                    "rd": 0,
                    "rs1": branch_gpr_index,
                    "simm12": 0,
                    "NoRestriction": 1,
                },
            )

            self._mHelperGprSet.releaseHelperGprs()

        return True


# This class generates instructions to update the system State according to
# the floating point register StateElements provided to the
# processStateElement() and processStateElements() methods.
class FloatingPointRegisterStateTransitionHandlerGPGPU(StateTransitionHandler):
    def __init__(self, aGenThread):
        super().__init__(aGenThread)

        self._mHelperGprSet = StateTransitionHelperGprSet(self)

    # Execute the State change represented by the StateElement. Only instances
    # of the StateElement types for which the StateTransitionHandler has been
    # registered will be passed to this method. Other StateTransitionHandlers
    # will process the other StateElement types. It is important to avoid
    # making changes to entities represented by StateElements that have already
    # been processed. Changes to entities represented by StateElements that
    # will be processed later are permitted.
    #
    #  @param aStateElem A StateElement object.
    def processStateElement(self, aStateElem):
        if aStateElem.getStateElementType() != EStateElementType.FloatingPointRegister:
            return False

        (reg_val_gpr_index,) = self._mHelperGprSet.acquireHelperGprs(1)

        load_gpr64_seq = LoadGPR64(self.genThread)
        load_gpr64_seq.load(reg_val_gpr_index, aStateElem.getValues()[0])

        (mem_block_ptr_index,) = self._mHelperGprSet.acquireHelperGprs(1)
        self.initializeMemoryBlock(mem_block_ptr_index, [aStateElem])
        fp_load_instr = "FLW##CALM" if aStateElem.getName().startswith("S") else "FLD##CALM"
        self.genInstruction(
            fp_load_instr,
            {
                "rd": aStateElem.getRegisterIndex(),
                "rs1": mem_block_ptr_index,
                "simm12": 0,
                "NoRestriction": 1,
            },
        )

        self._mHelperGprSet.releaseHelperGprs()

        return True

    # Execute the State changes represented by the StateElements. Only
    # instances of the StateElement types for which the StateTransitionHandler
    # has been registered will be passed to this method. Other
    # StateTransitionHandlers will process the other StateElement types. It is
    # important to avoid making changes to entities represented by
    # StateElements that have already been processed. Changes to entities
    # represented by StateElements that will be processed later are permitted.
    #
    #  @param aStateElems A list of all StateElement objects of a particular
    #       type.
    def processStateElements(self, aStateElems):
        if aStateElems[0].getStateElementType() != EStateElementType.FloatingPointRegister:
            self.error(
                "This StateTransitionHandler can only process "
                "StateElements of type %s" % EStateElementType.FloatingPointRegister
            )

        # The logic below assumes there will never be more than 2,048 floating
        # point register StateElements
        if len(aStateElems) > 2048:
            self.error(
                "Unexpected number of floating point register "
                "StateElements: %d" % len(aStateElems)
            )

        (mem_block_ptr_index,) = self._mHelperGprSet.acquireHelperGprs(1)
        self.initializeMemoryBlock(mem_block_ptr_index, aStateElems)
        offset = 0
        for state_elem in aStateElems:
            fp_load_instr = "FLW##CALM" if state_elem.getName().startswith("S") else "FLD##CALM"
            self.genInstruction(
                fp_load_instr,
                {
                    "rd": state_elem.getRegisterIndex(),
                    "rs1": mem_block_ptr_index,
                    "simm12": offset,
                    "NoRestriction": 1,
                },
            )
            offset += 8

        self._mHelperGprSet.releaseHelperGprs()


# This class assists with allocating temporary GPRs for use in
# StateTransitionHandler methods. It uses arbitrary GPRs where possible.
# However, it can also return non-arbitrary GPRs and reset those GPRs when they
# are no longer needed, avoiding any permanent State modifications.
class StateTransitionHelperGprSet(object):
    def __init__(self, aStateTransHandler):
        self._mStateTransHandler = aStateTransHandler
        self._mArbitraryGprIndices = []
        self._mNonArbitraryGprIndices = []
        self._mNonArbitraryGprOrigValues = []

    # Return GPRs that can be used to help effect a StateTransition. Arbitrary
    # GPRs are returned when possible. This call must be followed by a call to
    # releaseHelperGprs().
    #
    #  @param aGprCount The number of GPRs requested.
    #  @param aValidate Verify the StateTransition type supports using
    #       non-arbitrary GPRs.
    def acquireHelperGprs(self, aGprCount, aValidate=True):
        if self._mNonArbitraryGprIndices:
            self.error(
                "acquireHelperGprs() cannot be called without a matching call "
                "to releaseHelperGprs()"
            )

        self._mArbitraryGprIndices = self._mStateTransHandler.getArbitraryGprs(
            aGprCount, aExclude=(0, 1, 2)
        )
        if self._mArbitraryGprIndices is None:
            if aValidate:
                self._validateInsufficientArbitaryGprs(aGprCount)
            handler = self._mStateTransHandler
            self._mArbitraryGprIndices = handler.getAllArbitraryGprs(aExclude=(0, 1, 2))
            remaining_gpr_count = aGprCount - len(self._mArbitraryGprIndices)
            excluded_regs = "0,1,2,%s" % ",".join(
                str(index) for index in self._mArbitraryGprIndices
            )
            self._mNonArbitraryGprIndices = handler.getRandomGPRs(
                remaining_gpr_count, exclude=excluded_regs, no_skip=True
            )

        for non_arbitrary_gpr_index in self._mNonArbitraryGprIndices:
            (gpr_val, _) = self._mStateTransHandler.readRegister("x%d" % non_arbitrary_gpr_index)
            self._mNonArbitraryGprOrigValues.append(gpr_val)

        helper_gpr_indices = list(self._mArbitraryGprIndices)
        helper_gpr_indices.extend(self._mNonArbitraryGprIndices)
        return helper_gpr_indices

    # Indicate any acquired helper GPRs are no longer needed. Instructions are
    # generated to reset the values of any non-arbitrary GPRs.
    def releaseHelperGprs(self):
        load_gpr64_seq = LoadGPR64(self._mStateTransHandler.genThread)
        for (i, non_arbitrary_gpr_index) in enumerate(self._mNonArbitraryGprIndices):
            load_gpr64_seq.load(non_arbitrary_gpr_index, self._mNonArbitraryGprOrigValues[i])

        self._mNonArbitraryGprOrigValues = []
        self._mNonArbitraryGprIndices = []
        self._mArbitraryGprIndices = []

    # Fail if the StateTransition type does not support proceeding without the
    # requested number arbitrary GPRs available.
    #
    #  @param aGprCount The number of GPRs requested.
    def _validateInsufficientArbitaryGprs(self, aGprCount):
        if self._mStateTransHandler.mStateTransType == EStateTransitionType.Boot:
            self._mStateTransHandler.error(
                "Boot StateTransitions must have at least %d arbitary GPR(s) "
                "available" % aGprCount
            )


#  Return the value of a register with the specified field value inserted at
#  the appropriate place.
#
#  @param aSequence A Sequence object.
#  @param aRegName The name of the register.
#  @param aRegVal The current value of the register.
#  @param aRegFieldName The name of the register field.
#  @param aRegFieldVal The desired value of the register field.
def combine_register_value_with_field_value(
    aSequence, aRegName, aRegVal, aRegFieldName, aRegFieldVal
):
    (reg_field_mask, reg_field_reverse_mask) = aSequence.getRegisterFieldMask(
        aRegName, [aRegFieldName]
    )
    reg_field_shift = UtilityFunctions.lowestBitSet(reg_field_mask)
    reg_field_val = (aRegFieldVal << reg_field_shift) & reg_field_mask
    reg_val = (aRegVal & reg_field_reverse_mask) | reg_field_val

    return reg_val
